/*
 * MultirootSolver.cpp
 *
 *  Created on: 23 Jun 2011
 *      Author: allan
 */

#include "MultirootSolver.h"

// C++ includes
#include <iostream>
using namespace std;

// GSL includes
#include <gsl/gsl_multiroots.h>

int GSL_Function(const gsl_vector* x, void* params, gsl_vector* F);
int GSL_Jacobian(const gsl_vector* x, void* params, gsl_matrix* J);
int GSL_FunctionJacobian(const gsl_vector* x, void* params, gsl_vector* F, gsl_matrix* J);

MultirootSolver::MultirootSolver() :
type(MULTIROOTSOLVER_GNEWTON), tolerance(1.0E-6), maxiterations(100)
{}

void MultirootSolver::SetSolverType(MultirootSolverType type)
{
	this->type = type;
}

void MultirootSolver::SetSolverOptions(double tolerance, unsigned maxiterations)
{
	this->tolerance     = tolerance;
	this->maxiterations = maxiterations;
}

bool MultirootSolver::Solve(MultirootProblem& nonlinearProblem, VectorXd& x)
{
	// The number of equations in the non-linear problem
	const unsigned numEquations = x.size();
	
	//==========================================================================================//
	// Initialise the GSL function structure (gsl_func)
	//==========================================================================================//
	gsl_multiroot_function_fdf gsl_func;
	
	gsl_func.n      = numEquations;
	gsl_func.f      = &GSL_Function;
	gsl_func.df     = &GSL_Jacobian;
	gsl_func.fdf    = &GSL_FunctionJacobian;
	gsl_func.params = &nonlinearProblem;
	
	//==========================================================================================//
	// Initialise the GSL multiroot solver (gsl_solver)
	//==========================================================================================//
	gsl_multiroot_fdfsolver* gsl_solver;
	
	switch(type)
	{
	case MULTIROOTSOLVER_HYBRIDSJ: 
		gsl_solver = gsl_multiroot_fdfsolver_alloc(gsl_multiroot_fdfsolver_hybridsj, numEquations); 
		break;
	case MULTIROOTSOLVER_HYBRIDJ:
		gsl_solver = gsl_multiroot_fdfsolver_alloc(gsl_multiroot_fdfsolver_hybridj, numEquations); 
		break;
	case MULTIROOTSOLVER_NEWTON:
		gsl_solver = gsl_multiroot_fdfsolver_alloc(gsl_multiroot_fdfsolver_newton, numEquations);
		break;
	case MULTIROOTSOLVER_GNEWTON:
		gsl_solver = gsl_multiroot_fdfsolver_alloc(gsl_multiroot_fdfsolver_gnewton, numEquations);
		break;
	default:
		gsl_solver = gsl_multiroot_fdfsolver_alloc(gsl_multiroot_fdfsolver_gnewton, numEquations);
		break;
	}
	
	//==========================================================================================//
	// Initialise the GSL vector (gsl_x)
	//==========================================================================================//
	// Allocate memory for the GSL vector gsl_x
	gsl_vector* gsl_x = gsl_vector_alloc(numEquations);
	
	// Transfer the data in the Eigen vector x to the GSL vector gsl_x
	for(unsigned i = 0; i < numEquations; ++i)
		gsl_vector_set(gsl_x, i, x[i]);
	
	//==========================================================================================//
	// Register the GSL vector (gsl_x) and GSL function structure in the GSL solver (gsl_solver)
	//==========================================================================================//
	gsl_multiroot_fdfsolver_set(gsl_solver, &gsl_func, gsl_x);
	
	//==========================================================================================//
	// Begin the iteration procedure 
	//==========================================================================================//
	// A status variable
	int status = GSL_CONTINUE;
	
	// An iteration counter variable 
	unsigned counter = 0;
	
	// Start the iterations...
	do
	{
		// Iterate and get the status
		status = gsl_multiroot_fdfsolver_iterate(gsl_solver);
		
		// Check is the iteration was well // TODO It is necessary an error feedback here
		if(status) { cerr << "Error: The iteration " << counter << " in the MultirooSolver did not succeed," << endl; break; }
		
		// Transfer the data in the GSL vector (gsl_x) to the Eigen vector x
		for(unsigned i = 0; i < numEquations; ++i)
			x[i] = gsl_vector_get(gsl_solver->x, i);
	
		// Apply the post-processing method of the non-linear problem object in the solution vector
		nonlinearProblem.PostProcess(x);
		
		// Check the residual of the solution in terms of the function evaluation
		status = gsl_multiroot_test_residual(gsl_solver->f, tolerance);
		
		// Increment counter
		++counter;
	}
	while(status == GSL_CONTINUE && counter < maxiterations);

	// Transfer the data in the GSL vector (gsl_x) to the Eigen vector x
	for(unsigned i = 0; i < numEquations; ++i)
		x[i] = gsl_vector_get(gsl_solver->x, i);
	
	// Free allocated memory for the GSL data
	gsl_vector_free(gsl_x);
	gsl_multiroot_fdfsolver_free(gsl_solver);
	
	// Check if the method converged
	bool converged = counter < maxiterations;
	
	return converged; 
}

int GSL_Function(const gsl_vector* gslx, void* params, gsl_vector* gslF)
{
	MultirootProblem& problem = *((MultirootProblem*)params);
	
	// Create Eigen vectors x and F
	VectorXd x(gslx->size);
	VectorXd F(gslx->size);
	
	// Transfer the data in the GSL vector gslx to the Eigen vector x
	for(unsigned i = 0; i < gslx->size; ++i)
		x[i] = gsl_vector_get(gslx, i);
	
	// Compute F using the multiroot problem object
	problem.Function(x, F);
	
	// Transfer the data in the Eigen vector F to the GSL vector gslF
	for(unsigned i = 0; i < gslx->size; ++i)
		gsl_vector_set(gslF, i, F[i]); // TODO we could use copy algorithms here, since Eigen objects allow to return a pointer to double* (its data)
	
	return GSL_SUCCESS;
}

int GSL_Jacobian(const gsl_vector* gslx, void* params, gsl_matrix* gslJ)
{
	MultirootProblem& problem = *((MultirootProblem*)params);
	
	// Create an Eigen vector x and a matrix J
	VectorXd x(gslx->size);
	MatrixXd J(gslx->size, gslx->size);
	
	// Transfer the data in the GSL vector gslx to the Eigen vector x
	for(unsigned i = 0; i < gslx->size; ++i)
		x[i] = gsl_vector_get(gslx, i);
	
	// Compute J using the multiroot problem object
	problem.Jacobian(x, J);
	
	// Transfer the data in the Eigen matrix J to the GSL matrix gslJ 
	for(unsigned i = 0; i < gslx->size; ++i)
		for(unsigned j = 0; j < gslx->size; ++j)
			gsl_matrix_set(gslJ, i, j, J(i, j)); // TODO we could use copy algorithms here, since Eigen objects allow to return a pointer to double* (its data)
	
	return GSL_SUCCESS;
}

int GSL_FunctionJacobian(const gsl_vector* gslx, void* params, gsl_vector* gslF, gsl_matrix* gslJ)
{
	MultirootProblem& problem = *((MultirootProblem*)params);
	
	// Initialiase an Eigen vector x with the data in the GSL vector gslx
	VectorXd x(gslx->size);
	VectorXd F(gslx->size);
	MatrixXd J(gslx->size, gslx->size);
	
	// Transfer the data in the GSL vector gslx to the Eigen vector x
	for(unsigned i = 0; i < gslx->size; ++i)
		x[i] = gsl_vector_get(gslx, i);
	
	// Compute F and J using the multiroot problem object
	problem.FunctionJacobian(x, F, J);
	
	// Transfer the data in the Eigen vector F to the GSL vector gslF
	for(unsigned i = 0; i < gslx->size; ++i)
		gsl_vector_set(gslF, i, F[i]); // TODO we could use copy algorithms here, since Eigen objects allow to return a pointer to double* (its data)
	
	// Transfer the data in the Eigen matrix J to the GSL matrix gslJ 
	for(unsigned i = 0; i < gslx->size; ++i)
		for(unsigned j = 0; j < gslx->size; ++j)
			gsl_matrix_set(gslJ, i, j, J(i, j)); // TODO we could use copy algorithms here, since Eigen objects allow to return a pointer to double* (its data)
	
	return GSL_SUCCESS;
}
